# -*- coding: utf-8 -*-

import torch.nn as nn

from .token import TokenEmbedding
from .segment import SegmentEmbedding
from .position import PositionalEmbedding


class BERTEmbedding(nn.Module):
    """
    BERT的嵌入由以下特征组成：
        1. TokenEmbedding :         normal embedding matrix
        2. PositionalEmbedding :    使用sin，cos添加位置信息
        2. SegmentEmbedding :       添加句段信息，(sent_A:1, sent_B:2)

        所有这些功能的总和是BERTEmbedding的输出
    """

    def __init__(self, vocab_size, embed_size, dropout=0.1):
        """
        :param vocab_size:  总词典大小
        :param embed_size:  token embedding 的嵌入大小
        :param dropout:     失活率
        """
        super().__init__()
        self.token = TokenEmbedding(vocab_size=vocab_size, embed_size=embed_size)
        self.position = PositionalEmbedding(d_model=self.token.embedding_dim)
        self.segment = SegmentEmbedding(embed_size=self.token.embedding_dim)
        self.dropout = nn.Dropout(p=dropout)    # embedding嵌入失活率
        self.embed_size = embed_size

    def forward(self, sequence, segment_label):
        """
        bert的embedding组合过程
        :param sequence:
        :param segment_label:
        :return:
        """
        x = self.token(sequence) + self.position(sequence) + self.segment(segment_label)
        return self.dropout(x)
